#!/usr/bin/env python
import rospy
import numpy as np
import open3d as o3d
import ros_np_multiarray as ros_nm
from vision_messages.srv import GraspPoseEst, GraspPoseEstRequest, GraspPoseEstResponse
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Int64
from geometry_msgs.msg import Point
from vision_messages.msg import SingleObjPose

def xyz_array_to_pointcloud2(points, stamp=None, frame_id=None):
    '''
    Create a sensor_msgs.PointCloud2 from an array
    of points.
    '''
    msg = PointCloud2()
    if stamp:
        msg.header.stamp = stamp
    if frame_id:
        msg.header.frame_id = frame_id
    if len(points.shape) == 3:
        msg.height = points.shape[1]
        msg.width = points.shape[0]
    else:
        msg.height = 1
        msg.width = len(points)
    msg.fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1)]
    msg.is_bigendian = False
    msg.point_step = 12
    msg.row_step = 12*points.shape[0]
    msg.is_dense = int(np.isfinite(points).all())
    msg.data = np.asarray(points, np.float32).tostring()
    return msg

def read_pcd(file_path):
    pcd = o3d.io.read_point_cloud(file_path)
    print(np.asarray(pcd.points))
    colors = np.asarray(pcd.colors) * 255
    points = np.asarray(pcd.points)
    print(points.shape, colors.shape)
    return points

def preprocess(data_in):
    if isinstance(data_in, np.ndarray):
        recon_req = SingleObjPose()
        recon_req.cam_pcs = ros_nm.to_multiarray_f32(data_in)
    else:
        raise TypeError('data_in type error!')

    return recon_req

# TODO: change Panda TO JACO
def postprocess(response):
    assert isinstance(response, GraspPoseEstResponse)
    g_RT = ros_nm.to_numpy_f32(response.grasp_pose)
    print(g_RT)
    return g_RT
    
class GraspNet_Client:
    def __init__(self):
        self.grasp_est_srv = rospy.ServiceProxy('/perception/graspnet', GraspPoseEst)
        rospy.wait_for_service('/perception/graspnet')
        
    def get_grasp_pose(self, sp_result):
        #################################
        # example in debug:
        #################################
        # req = detect_graspsRequest()
        # req.cloud_indexed = CloudIndexed()
        # req.cloud_indexed.cloud_sources = CloudSources()
        # pc = read_pcd(file_path='/home/robot/Work/gpd/tutorials/krylon.pcd')
        # req.cloud_indexed.cloud_sources.cloud = xyz_array_to_pointcloud2(pc)
        # req.cloud_indexed.cloud_sources.camera_source = list()
        # req.cloud_indexed.cloud_sources.view_points = list()
        # for i in range(0, len(pc)):
        #     req.cloud_indexed.cloud_sources.camera_source.append(Int64(data=0))
        # req.cloud_indexed.cloud_sources.view_points.append(Point()) # initilize to {x=0,y=0,z=0}
        # try:
        #     res = self.grasp_est_srv(req)
        #     print(res)
        #     return res
        # except rospy.ServiceException as e:            
        #     rospy.loginfo("Grasp Estimation Service call failed: %s"%e)
        ##################################

        ##################################
        req = preprocess(sp_result)
        try:
            res = self.grasp_est_srv(req)
            return postprocess(res)
        except rospy.ServiceException as e:            
            rospy.loginfo("Grasp Estimation Service call failed: %s"%e)
        

if __name__ == '__main__':
    rospy.init_node('GraspNet_client')
    grasp_est_client = GraspNet_Client()
    grasp_est_client.get_grasp_pose()